package tree;

import java.util.Comparator;

import tree.BinaryTree.Node;

public class BST<E> extends BinaryTree<E>{
	
	private Comparator<E> comparator;
	
	
	public BST() {
		this(null);
	}
	public BST(Comparator<E> comparator) {
		this.comparator = comparator;
	}
	
	
	public void add(E element) {
		checkElement(element);
		if (root == null) {
			root = createNode(element, null);
			size++;
			afterAdd(root);
			return;
		}
		Node<E> node = root;
		Node<E> parent = root;
		int cmp = 0;
		while (node != null) {
			parent = node;
			cmp = compare(element, node.element);
			if (cmp > 0) {
				node = node.right;
			} else if (cmp < 0) {
				node = node.left;
			} else {
				node.element = element;
				return;
			}
		}
		Node<E> newNode = createNode(element, parent);
		if (cmp > 0) {
			parent.right = newNode;
		} else {
			parent.left = newNode;
		}
		size++;
		afterAdd(newNode);
	}
	
	public void remove(E element) {
		remove(node(element));
	}
	
	private void remove(Node<E> node) {
		if (node == null) return;
		size--;
		if (node.hasTwoChildren()) {
			Node<E> sucNode = successor(node);
			node.element = sucNode.element;
			node = sucNode;
		}
		
		Node<E> replacement = node.left != null ? node.left : node.right;
		
		if (replacement != null) {
			replacement.parent = node.parent;
			if (node.parent == null) {
				root = replacement;
			} else if (node == node.parent.left) {
				node.parent.left = replacement;
			} else {
				node.parent.right = replacement;
			}
			afterRemove(replacement);
		} else if (node.parent == null) {
			root = null;
			afterRemove(node);
		} else {
			if (node == node.parent.left) {
				node.parent.left = null;
			} else {
				node.parent.right = null;
			}
			afterRemove(node);
		} 
		
		
	}
	
	public boolean contains(E element) {
		return node(element) != null;
	}
	
	protected Node<E> createNode(E element, Node<E> parent) {
		return new Node<>(element, parent);
	}
	
	/**
	 * 添加节点后对树的结构进行调整(子类实现)
	 * @param node
	 */
	protected void afterAdd(Node<E> node) {}
	/**
	 * 删除节点后对树的结构进行调整(子类实现)
	 * @param node
	 */
	protected void afterRemove(Node<E> node) {}
	
	/**
	 * 根据元素值查找节点
	 * @param element
	 * @return
	 */
	private Node<E> node(E element) {
		checkElement(element);
		Node<E> node = root;
		int cmp = 0;
		while (node != null) {
			cmp = compare(element, node.element);
			if (cmp == 0) {
				return node;
			} else if (cmp > 0) {
				node = node.right;
			} else {
				node = node.left;
			}
		}
		return null;
	}
	
	private int compare(E e1, E e2) {
		if (comparator != null) return comparator.compare(e1, e2);
		return ((Comparable)e1).compareTo(e2);
	}
	
}
